# -*- coding: utf-8 -*-
# time: 2025/4/9 11:42
# file: ch01.py
# author: hanson
from langchain.agents import create_tool_calling_agent, AgentExecutor
from langchain.tools import Tool
from langchain_core.prompts import ChatPromptTemplate
from langchain_ollama import ChatOllama

# 2. 创建提示模板
prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个数学计算助手"),  # 系统角色设定
    ("human", "{input}"),  # 用户输入
    ("placeholder", "{agent_scratchpad}") # 代理工作区（必须包含）
])
def calculate_sum(number: int) -> str:
    """计算水果总数"""
    return str(number) # 4个桃子

math_tool = Tool(
    name="FruitCalculator",
    func=calculate_sum,
    description="计算水果总数，输入苹果数量，返回苹果+桃子的总和",
    return_direct=True,
)
# 创建LLM
llm = ChatOllama(model="qwen2.5:1.5b", temperature=0)
# 创建代理 和 工具
tools = [math_tool]
agent = create_tool_calling_agent(llm, tools,prompt)
# 执行代理
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

result = agent_executor.invoke({
    "input": "桌上有3个苹果和4个桃子，共有几个水果？"
})

print(result)
print(result["output"])